import numpy as np
import sys
from scipy.io import savemat
from scipy.sparse import coo_matrix, csc_matrix, csr_matrix, block_diag
from scipy.sparse import spdiags, eye, bmat, tril, triu
from scipy.sparse.linalg import cg, inv, dsolve,  gmres, minres, LinearOperator, spsolve_triangular
from scipy.sparse.linalg import spsolve
import pyamg

sys.path.append('../Boundarycondition')
sys.path.append('../solver')
sys.path.append('../functionspace')

from fealpy.functionspace.VectorLagrangeFiniteElementSpace import VectorLagrangeFiniteElementSpace
from fealpy.functionspace.LagrangeFiniteElementSpace import LagrangeFiniteElementSpace
from fealpy.fem.integral_alg import IntegralAlg
from fealpy.fem.doperator import stiff_matrix
from boundarycondition import NatureBC
from boundarycondition import EssentialBC
from Inter_Solver import Inter_Solver
from timeit import default_timer as timer
import matlab.engine
import cProfile

class LinearElasticityFEMModel:
    def __init__(self, mesh,  pde, p, integrator):
        dim = mesh.geo_dimension()
        if dim == 2:
            from HuZhangFiniteElementSpace2D import HuZhangFiniteElementSpace
        elif dim == 3:
            from HuZhangFiniteElementSpace3D import HuZhangFiniteElementSpace

        self.mesh = mesh
        self.tensorspace = HuZhangFiniteElementSpace(mesh, p)
        self.vectorspace = VectorLagrangeFiniteElementSpace(mesh, p-1, spacetype='D') 
        self.cvectorspace = VectorLagrangeFiniteElementSpace(mesh,1) #linear vector space
        self.cspace = LagrangeFiniteElementSpace(mesh, 1) # linear space 
        self.dim = self.tensorspace.dim
        self.sh = self.tensorspace.function()
        self.uh = self.vectorspace.function()
        self.pde = pde
        self.x = None
        self.integrator = integrator
        self.measure = mesh.entity_measure()
        self.integralalg = IntegralAlg(self.integrator, mesh, self.measure)
        self.count = 0

        self.itype = mesh.itype
        self.ftype = mesh.ftype

    def get_cvector_matrix(self):
        cvectorspace = self.cvectorspace
        bcs, ws = self.integrator.quadpts, self.integrator.weights
        cellmeasure = self.measure

        grad_basis = cvectorspace.grad_basis(bcs)
        eplison_basis = grad_basis
        eplison_basis[...,0,1] = 0.5*(eplison_basis[...,0,1]+grad_basis[...,1,0]) #(NQ,NC,ldof,gdim,gdim)
        eplison_basis[...,0,1] = 0.5*(eplison_basis[...,1,0]+grad_basis[...,0,1]) #(NQ,NC,ldof,gdim,gdim)

        A = 2*self.pde.mu*np.einsum('i,ijklm,ijslm,j->jks',ws,eplison_basis,eplison_basis,cellmeasure) #(NC,ldof,ldof)
        cell2dof = cvectorspace.cell_to_dof() #(NC,ldof)

        I = np.broadcast_to(cell2dof[:,:,None],shape=A.shape)
        J = np.broadcast_to(cell2dof[:,None,:],shape=A.shape)

        gdof = cvectorspace.number_of_global_dofs()

        A = csr_matrix((A.flat,(I.flat,J.flat)),shape=(gdof,gdof))

        isBdDof = cvectorspace.boundary_dof_flag()

        bdIdx = np.zeros((A.shape[0], ), np.int)
        bdIdx[isBdDof] = 1
        Tbd = spdiags(bdIdx, 0, A.shape[0], A.shape[0])
        T = spdiags(1-bdIdx, 0, A.shape[0], A.shape[0])
        A = T@A@T + Tbd
        return A




        

    def precondieitoner(self):

        tspace = self.tensorspace
        vspace = self.vectorspace
        tgdof = tspace.number_of_global_dofs()

        gdim = tspace.geo_dimension()

        bcs, ws = self.integrator.quadpts, self.integrator.weights
        phi = tspace.basis(bcs)

        # construct diag matrix D
        if gdim == 2:
            d = np.array([1, 1, 2])
        elif gdim == 3:
            d = np.array([1, 1, 1, 2, 2, 2])
        #print(phi.shape) #(NQ,NC,ldof,tdim,tdim)
        D = np.einsum('i, ijklm, m, ijklm, j->jkl', ws, phi, d, phi, self.measure) #(NC,ldof,tdim)
        tcell2dof = tspace.cell_to_dof() #(NC,ldof,tdim)
        self.D = np.bincount(tcell2dof.flat, weights=D.flat, minlength=tgdof) 
        self.D[self.isBdDof] = 1.0

        # construct amg solver 
        A = stiff_matrix(self.cspace, self.integrator, self.measure)
        isBdDof = self.cspace.boundary_dof()
        bdIdx = np.zeros((A.shape[0], ), np.int)
        bdIdx[isBdDof] = 1
        #bdIdx[0] = 1
        Tbd = spdiags(bdIdx, 0, A.shape[0], A.shape[0])
        T = spdiags(1-bdIdx, 0, A.shape[0], A.shape[0])
        A = T@A@T + Tbd
        self.A = A
        if gdim == 2:
            self.ml = pyamg.ruge_stuben_solver(A) # 这里要求必须有网格内部节点 

        # Get interpolation matrix 
        NC = self.mesh.number_of_cells()
        bc = self.vectorspace.dof.multiIndex/self.vectorspace.p #(ldof,gdim+1)
        val = np.tile(bc, (NC, 1)) #(NC*ldof,gdim+1)
        #print(val.shape,bc.shape,NC)
        c2d0 = self.vectorspace.dof.cell2dof #(NC,ldof)
        c2d1 = self.cspace.cell_to_dof() #(NC,cldof)
        #print(c2d0.shape,c2d1.shape)

        gdim = self.tensorspace.geo_dimension()
        I = np.einsum('ij, k->ijk', c2d0, np.ones(gdim+1,dtype=int)) #(NC,ldof,gdim+1)
        J = np.einsum('ik, j->ijk', c2d1, np.ones(len(bc),dtype=int)) #(NC,ldof,cldof)
        cgdof = self.cspace.number_of_global_dofs()
        fgdof = self.vectorspace.number_of_global_dofs()//self.mesh.geo_dimension()
        #print(val.shape,I.shape,J.shape,bc.shape)
        self.PI = csr_matrix((val.flat, (I.flat, J.flat)), shape=(fgdof, cgdof))
        #self.PI[:,isBdDof] = 0*self.PI[:,isBdDof]

    

    def precondieitoner_1(self):
        tspace = self.tensorspace
        vspace = self.vectorspace
        tgdof = tspace.number_of_global_dofs()

        gdim = tspace.geo_dimension()

        bcs, ws = self.integrator.quadpts, self.integrator.weights
        phi = tspace.basis(bcs)

        # construct diag matrix D
        if gdim == 2:
            d = np.array([1, 1, 2])
        elif gdim == 3:
            d = np.array([1, 1, 1, 2, 2, 2])
        #print(phi.shape) #(NQ,NC,ldof,tdim,tdim)
        D = np.einsum('i, ijklm, m, ijklm, j->jkl', ws, phi, d, phi, self.measure) #(NC,ldof,tdim)
        tcell2dof = tspace.cell_to_dof() #(NC,ldof,tdim)
        self.D = np.bincount(tcell2dof.flat, weights=D.flat, minlength=tgdof) 

        # construct amg solver 
        A = self.get_cvector_matrix()
        isBdDof = self.cvectorspace.boundary_dof_flag()
        bdIdx = np.zeros((A.shape[0], ), np.int)
        bdIdx[isBdDof] = 1
        Tbd = spdiags(bdIdx, 0, A.shape[0], A.shape[0])
        T = spdiags(1-bdIdx, 0, A.shape[0], A.shape[0])
        A = T@A@T + Tbd
        self.A = A
        if gdim == 2:
            self.ml = pyamg.ruge_stuben_solver(A) # 这里要求必须有网格内部节点

        # Get interpolation matrix 
        NC = self.mesh.number_of_cells()
        bc = self.vectorspace.dof.multiIndex/self.vectorspace.p #(ldof,gdim+1)
        ldof = len(bc)

        val = np.tile(bc,(NC,gdim)) #(NC*ldof,gdim*(gdim+1))

        c2d0 = self.vectorspace.cell_to_dof().reshape((NC,ldof,gdim))   #(NC,ldof,gdim)
        c2d1 = self.cvectorspace.cell_to_dof().reshape((NC,gdim+1,gdim))    #(NC,(gdim+1),gdim)

        I = np.einsum('ijk,l->ijlk',c2d0,np.ones(gdim+1,dtype=int)) #(NC,ldof,gdim+1,gdim)
        J = np.einsum('ilk,j->ijlk',c2d1,np.ones(ldof,dtype=int)) #(NC,ldof,gdim+1,gdim)

        cgdof = self.cvectorspace.number_of_global_dofs()
        fgdof = self.vectorspace.number_of_global_dofs()

        self.PI = csr_matrix((val.flat, (I.flat, J.flat)), shape=(fgdof, cgdof))
        self.PI[:,isBdDof] = 0*self.PI[:,isBdDof]














    def get_left_matrix(self):
        mesh = self.mesh
        tspace = self.tensorspace
        vspace = self.vectorspace
        gdim = tspace.geo_dimension()
        bcs, ws = self.integrator.quadpts, self.integrator.weights
        
        if gdim == 2:
            d = np.array([1, 1, 2])
        elif gdim == 3:
            d = np.array([1, 1, 1, 2, 2, 2])

        tldof = tspace.number_of_local_dofs()
        tdim = tspace.tensor_dimension()
        NC = self.mesh.number_of_cells()
        NQ = bcs.shape[0]

        phi = tspace.basis(bcs).reshape(NQ,NC,-1,tdim) #(NQ,NC,ldof,tdim)
        aphi = self.pde.compliance_tensor(phi)
        M = np.einsum('i, ijkm, m, ijom, j->jko', ws, aphi, d, phi, self.measure, optimize=True)

        tcell2dof = tspace.cell_to_dof().reshape(NC,-1)
        I = np.einsum('ij, k->ijk', tcell2dof, np.ones(tldof))
        J = I.swapaxes(-1, -2)
        tgdof = tspace.number_of_global_dofs()
        M = csr_matrix((M.flat, (I.flat, J.flat)), shape=(tgdof, tgdof))

        vgdof = vspace.number_of_global_dofs()
        vldof = vspace.number_of_local_dofs()

        dphi = tspace.div_basis(bcs).reshape(NQ,NC,-1,gdim) #(NQ, NC, ldof*tdim, gdim)
        uphi = vspace.basis(bcs) #(NQ,1,vldof,gdim)
        #print(dphi.shape,uphi.shape)
        B = np.einsum('i, ijkm, ijom, j->jko', ws, uphi, dphi, self.measure, optimize=True)

        I = np.einsum('ij, k->ijk', vspace.cell_to_dof(), np.ones(tldof,dtype=int))
        J = np.einsum('ij, k->ikj', tspace.cell_to_dof().reshape(NC,-1,), np.ones(vldof,dtype=int))
        #print(B.shape,I.dtype,J.shape)
        B = csr_matrix((B.flat, (I.flat, J.flat)), shape=(vgdof, tgdof))
        return  M, B

    def get_right_vector(self):
        tspace = self.tensorspace
        vspace = self.vectorspace

        bcs, ws = self.integrator.quadpts, self.integrator.weights
        pp = vspace.mesh.bc_to_point(bcs)
        fval = self.pde.source(pp)#(NQ,NC,gdim)
        phi = vspace.basis(bcs)#(NQ,1,vldof,gdim)
        #print(phi.shape,fval.shape)
        bb = np.einsum('i, ikm, ikjm, k->kj', ws, fval, phi, self.measure)

        cell2dof = vspace.cell_to_dof()
        vgdof = vspace.number_of_global_dofs()
        b = np.bincount(cell2dof.flat, weights=bb.flat, minlength=vgdof)
        return  -b

    def get_AA_and_bb(self):
        pde = self.pde
        tensorspace = self.tensorspace
        vectorspace = self.vectorspace
        
        tgdof = self.tensorspace.number_of_global_dofs()
        vgdof = self.vectorspace.number_of_global_dofs()
        gdof = tgdof + vgdof

        start = timer()
        M, B = self.get_left_matrix()
        bb =  self.get_right_vector()
        AA = bmat([[M, B.transpose()], [B, None]]).tocsr()
        bb = np.r_[np.zeros(tgdof), bb]

        #print(np.max(np.abs(A-A.T)),bb), 此时A,b,对称        
        #边界给的是法向或者切向分量，完全体现再gD中
        bc = NatureBC(tensorspace,pde.dirichlet,threshold=pde.is_dirichlet_boundary)
        bb = bc.apply(bb)
        #print(np.max(np.abs(A-A.T)),bb)   ##此时A,b,对称 

        bc = EssentialBC(tensorspace,pde.neumann,threshold=pde.is_neumann_boundary)
        AA, bb = bc.apply(AA,bb)
        self.isBdDof = bc.isBdDof

        end = timer()
        print("Construct linear system time:", end - start)

        self.AA = AA
        self.bb = bb
        self.M = M
        self.B = B


    def solve(self):    
        idx = 0
        start = timer()
        ISr = Inter_Solver(self.AA,self.bb,maxiter=100000,tol=1e-15)
        if idx==0:
            x = spsolve(self.AA, self.bb)
        elif idx==1:
            x = ISr.MINRES()
            r = A@x-bb
            print(np.sqrt(r@r))
        elif idx==2:
            xx, exitCode = minres(self.AA,self.bb,maxiter=10000,tol=1e-15)
            r = A@xx-bb
            print(r@r)

        end = timer()
        print("Solve time:", end-start)
        self.x = x



    def fast_solve(self):
        pde = self.pde
        tensorspace = self.tensorspace
        vectorspace = self.vectorspace

        start = timer()
        #self.precondieitoner_1()
        self.precondieitoner()
        end = timer()
        print("Construct precondieitoner time:", end - start)

        ########################################################
        ########################################################

        tgdof = self.tensorspace.number_of_global_dofs()
        vgdof = self.vectorspace.number_of_global_dofs()
        gdof = tgdof + vgdof

      
        self.M, self.B = self.get_left_matrix()
        S = self.B@spdiags(1/self.D, 0, tgdof, tgdof)@self.B.transpose()
        self.SL = tril(S).tocsc()
        self.SU = triu(S, k=1).tocsr()

        self.SUT = self.SL.transpose().tocsr()
        self.SLT = self.SU.transpose().tocsr()



        self.internumber = 0 #调用预条件次数
        self.linear_operator_time = 0 #预条件时间
        start = timer()
        P = LinearOperator((gdof, gdof), matvec=self.linear_operator)

        x, exitCode = gmres(self.AA, self.bb, M=P, maxiter=10, tol=1e-8)
        r = self.AA@x-self.bb
        #print(exitCode,np.sqrt(r@r)/np.sqrt(self.bb@self.bb))
        end = timer()
        self.solve_time = end-start
        print("Fast solve time:", end-start)
        print("Inter number:",self.internumber)
        self.x = x
        


    def Finite_emelent_solution(self,x=None):
        tgdof = self.tensorspace.number_of_global_dofs()
        if x is not None:
            self.x = x

        if self.x is None:
            raise ValueError('Need to get x!')
        else:
            self.sh[:] = self.x[0:tgdof]
            self.uh[:] = self.x[tgdof:]
            

    
    def Save_A_and_b_to_matlab(self):
        AA = csc_matrix(self.AA)
        start = timer()
        savemat('/Users/chen/git/mixed-hzelem/test/data/data.mat',{'AA':AA,'bb':self.bb})
        end = timer()
        print("Save data time:", end - start)

    def Save_A_and_b_to_matlab_fast(self):
        start = timer()
        AA = csc_matrix(self.AA)
        PI = csc_matrix(self.PI)
        B = csc_matrix(self.B)
        A = csc_matrix(self.A)
        D = csc_matrix(self.D)

        tgdof = self.tensorspace.number_of_global_dofs()
        vgdof = self.vectorspace.number_of_global_dofs()
        gdim = self.mesh.geo_dimension()

        restart = 20
        tol = 1e-8
        maxit = 30

        S = csc_matrix(self.B@spdiags(1/self.D, 0, tgdof, tgdof)@self.B.transpose())
        savemat('/Users/chen/git/mixed-hzelem/test/data/fast_data.mat',{'AA':AA,'bb':self.bb,
                'D':D,'PI':PI,'S':S,'B':B,'A':A,'tgdof':tgdof,'vgdof':vgdof,'gdim':gdim,
                'restart':restart,'tol':tol,'maxit':maxit,'elem':self.mesh.entity('cell')})
        end = timer()






    def matlab_solve(self):
        self.Save_A_and_b_to_matlab()
        eng = matlab.engine.start_matlab()
        eng.addpath('/Users/chen/git/mixed-hzelem/test/solver')
        self.x = np.array(eng.matlab_solve())[:,0];
        
        





    def linear_operator(self, r):
        self.internumber+=1
        start = timer()

        tgdof = self.tensorspace.number_of_global_dofs()
        vgdof = self.vectorspace.number_of_global_dofs()
        gdof = tgdof + vgdof
        gdim = self.tensorspace.geo_dimension()

        r0 = r[0:tgdof]
        r1 = r[tgdof:]

        u0 = r0/self.D
        u1 = np.zeros(vgdof, dtype=np.float)
        r2 = r1 - self.B@u0

        for i in range(3):
            u1[:] = spsolve(self.SL, r2 - self.SU@u1, permc_spec="NATURAL") 

        r3 = r2 - (self.SL@u1 + self.SU@u1)

        if self.PI.shape[1] == self.cvectorspace.number_of_global_dofs():
            for i in range(gdim):
                if gdim == 2:
                    u1 = self.PI@self.ml.solve(self.PI.transpose()@r3, tol=1e-8, accel='cg')
                else:
                    u1 = self.PI@spsolve(self.A,self.PI.transpose()@r3,permc_spec="NATURAL")


        else:
            for i in range(gdim):
                if gdim == 2:
                    u1[i::gdim] = self.PI@self.ml.solve(self.PI.transpose()@r3[i::gdim], tol=1e-8, accel='cg')
                else:
                    u1[i::gdim] = self.PI@spsolve(self.A,self.PI.transpose()@r3[i::gdim],permc_spec="NATURAL")

        
        start = timer()
        for i in range(3):
            u1[:] = spsolve(self.SUT, r2 - self.SLT@u1, permc_spec="NATURAL")

        end = timer()
        self.linear_operator_time+=(end-start)


        #print(np.r_[u0 + self.B.transpose()@u1, -u1].shape)

        return np.r_[u0 + self.B.transpose()@u1/self.D, -u1]

    def linear_operator_minres(self,r):
        tgdof = self.tensorspace.number_of_global_dofs()
        vgdof = self.vectorspace.number_of_global_dofs()
        gdof = tgdof + vgdof
        gdim = self.tensorspace.geo_dimension()
        r0 = r[0:tgdof]
        r1 = r[tgdof:]

        u0 = r0/self.D
        u1 = np.zeros(vgdof, dtype=np.float)

        for i in range(3):
            u1[:] = spsolve(self.SL, r1 - self.SU@u1, permc_spec="NATURAL") 

        r2 = r1 - (self.SL@u1 + self.SU@u1)
        for i in range(gdim):
            u1[i::gdim] = self.PI@self.ml.solve(self.PI.transpose()@r2[i::gdim], tol=1e-8, accel='cg')
            #u1[i::gdim] = self.PI@spsolve(self.A,self.PI.transpose()@r2[i::gdim],permc_spec="NATURAL")

        for i in range(3):
            u1[:] = spsolve(self.SUT, r1 - self.SLT@u1, permc_spec="NATURAL")

        return np.r_[u0,u1]







    def error(self):
        self.sI = self.tensorspace.interpolation(self.pde.stress)
        self.uI = self.vectorspace.interpolation(self.pde.displacement)

        dim = self.dim
        mesh = self.mesh

        s = self.pde.stress
        sh = self.sh.value 
        e0 = self.integralalg.L2_error(s, sh)

        ds = self.pde.div_stress
        dsh = self.sh.div_value
        e1 = self.integralalg.L2_error(ds, dsh)

        u = self.pde.displacement
        uh = self.uh.value
        e2 = self.integralalg.L2_error(u, uh)

        sI = self.sI.value 
        e3 = self.integralalg.L2_error(s, sI)

        dsI = self.sI.div_value
        e4 = self.integralalg.L2_error(ds, dsI)

        uI = self.uI.value
        e5 = self.integralalg.L2_error(u, uI)

        e6 = self.integralalg.L2_error_uI_uh(uI, uh)

        return e0, e1, e2, e3, e4, e5, e6
